# python <pdb_name> <ori_chains> <aim_chains>
# python example.pdb ABC EFG

import sys
pdb_name = sys.argv[1]
ori_chains = sys.argv[2].upper().strip()
to_chains = sys.argv[3].upper().strip()

ori_chains_list = list(ori_chains)
to_chains_list = list(to_chains)

if len(ori_chains_list) != len(to_chains):
    print("输入的’初始链名‘数量应该等于’结果链名‘")
    sys.exit(0)


pdb = open("./{}.pdb".format(pdb_name.split(".pdb")[0]),"r")
pdb_lines = [i for i in pdb]

pdb.close()
for i in range(len(pdb_lines)):
    if pdb_lines[i][0:4]  == 'ATOM':
        for j in range(len(ori_chains_list)):
            if pdb_lines[i][21] == ori_chains_list[j]:
                pdb_lines[i] = pdb_lines[i][0:21] + to_chains_list[j] + pdb_lines[i][22:-1] + pdb_lines[i][-1]

pdb = open("./{}_chain_changed.pdb".format(pdb_name),"w")
pdb.writelines(pdb_lines)
pdb.close()